{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Get Okubo Weis\n\n\\begin{align}OW = S_n^2 + S_s^2 + \\omega^2\\end{align}\n\nwith normal strain ($S_n$), shear strain ($S_s$) and vorticity ($\\omega$)\n\n\\begin{align}S_n = \\frac{\\partial u}{\\partial x} - \\frac{\\partial v}{\\partial y},\n    S_s = \\frac{\\partial v}{\\partial x} + \\frac{\\partial u}{\\partial y},\n    \\omega = \\frac{\\partial v}{\\partial x} - \\frac{\\partial u}{\\partial y}\\end{align}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\nfrom numpy import arange, ma, where\n\nfrom py_eddy_tracker import data\nfrom py_eddy_tracker.dataset.grid import RegularGridDataset\nfrom py_eddy_tracker.observations.observation import EddiesObservations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def start_axes(title, zoom=False):\n    fig = plt.figure(figsize=(12, 6))\n    axes = fig.add_axes([0.03, 0.03, 0.90, 0.94])\n    axes.set_xlim(0, 360), axes.set_ylim(-80, 80)\n    if zoom:\n        axes.set_xlim(270, 340), axes.set_ylim(20, 50)\n    axes.set_aspect(\"equal\")\n    axes.set_title(title)\n    return axes\n\n\ndef update_axes(axes, mappable=None):\n    axes.grid()\n    if mappable:\n        plt.colorbar(mappable, cax=axes.figure.add_axes([0.94, 0.05, 0.01, 0.9]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load detection files\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "a = EddiesObservations.load_file(data.get_demo_path(\"Anticyclonic_20190223.nc\"))\nc = EddiesObservations.load_file(data.get_demo_path(\"Cyclonic_20190223.nc\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load Input grid, ADT will be used to detect eddies\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "g = RegularGridDataset(\n    data.get_demo_path(\"nrt_global_allsat_phy_l4_20190223_20190226.nc\"),\n    \"longitude\",\n    \"latitude\",\n)\n\nax = start_axes(\"ADT (cm)\")\nm = g.display(ax, \"adt\", vmin=-120, vmax=120, factor=100)\nupdate_axes(ax, m)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Get parameter for ow\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "u_x = g.compute_stencil(g.grid(\"ugos\"))\nu_y = g.compute_stencil(g.grid(\"ugos\"), vertical=True)\nv_x = g.compute_stencil(g.grid(\"vgos\"))\nv_y = g.compute_stencil(g.grid(\"vgos\"), vertical=True)\now = g.vars[\"ow\"] = (u_x - v_y) ** 2 + (v_x + u_y) ** 2 - (v_x - u_y) ** 2\n\nax = start_axes(\"Okubo weis\")\nm = g.display(ax, \"ow\", vmin=-1e-10, vmax=1e-10, cmap=\"bwr\")\nupdate_axes(ax, m)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Gulf stream zoom\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ax = start_axes(\"Okubo weis, Gulf stream\", zoom=True)\nm = g.display(ax, \"ow\", vmin=-1e-10, vmax=1e-10, cmap=\"bwr\")\nkw_ed = dict(intern_only=True, color=\"k\", lw=1)\na.display(ax, **kw_ed), c.display(ax, **kw_ed)\nupdate_axes(ax, m)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "only negative OW\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ax = start_axes(\"Okubo weis, Gulf stream\", zoom=True)\nthreshold = ow.std() * -0.2\now = ma.array(ow, mask=ow > threshold)\nm = g.display(ax, ow, vmin=-1e-10, vmax=1e-10, cmap=\"bwr\")\na.display(ax, **kw_ed), c.display(ax, **kw_ed)\nupdate_axes(ax, m)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Get okubo-weiss mean/min/center in eddies\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 6))\nax = plt.subplot(111)\nax.set_xlabel(\"Okubo-Weiss parameter\")\nkw_hist = dict(bins=arange(-20e-10, 20e-10, 50e-12), histtype=\"step\")\nfor method in (\"mean\", \"center\", \"min\"):\n    kw_interp = dict(grid_object=g, varname=\"ow\", method=method, intern=True)\n    _, _, m = ax.hist(\n        a.interp_grid(**kw_interp), label=f\"Anticyclonic - OW {method}\", **kw_hist\n    )\n    ax.hist(\n        c.interp_grid(**kw_interp),\n        label=f\"Cyclonic - OW {method}\",\n        color=m[0].get_edgecolor(),\n        ls=\"--\",\n        **kw_hist,\n    )\nax.axvline(threshold, color=\"r\")\nax.set_yscale(\"log\")\nax.grid()\nax.set_ylim(1, 1e4)\nax.set_xlim(-15e-10, 15e-10)\nax.legend()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Catch eddies with bad OW\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ax = start_axes(\"Eddies with a min OW in speed contour over threshold\")\now_min = a.interp_grid(**kw_interp)\na_bad_ow = a.index(where(ow_min > threshold)[0])\na_bad_ow.display(ax, color=\"r\", label=\"Anticyclonic\")\now_min = c.interp_grid(**kw_interp)\nc_bad_ow = c.index(where(ow_min > threshold)[0])\nc_bad_ow.display(ax, color=\"b\", label=\"Cyclonic\")\nax.legend()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Display Radius and amplitude of eddies\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(figsize=(12, 5))\nfig.suptitle(\n    \"Parameter distribution (solid line) and cumulative distribution (dashed line)\"\n)\nax_amp, ax_rad = fig.add_subplot(121), fig.add_subplot(122)\nax_amp_c, ax_rad_c = ax_amp.twinx(), ax_rad.twinx()\nax_amp_c.set_ylim(0, 1), ax_rad_c.set_ylim(0, 1)\nkw_a = dict(xname=\"amplitude\", bins=arange(0, 2, 0.002).astype(\"f4\"))\nkw_r = dict(xname=\"radius_s\", bins=arange(0, 500e6, 2e3).astype(\"f4\"))\nfor d, label, color in (\n    (a, \"Anticyclonic all\", \"r\"),\n    (a_bad_ow, \"Anticyclonic bad OW\", \"orange\"),\n    (c, \"Cyclonic all\", \"blue\"),\n    (c_bad_ow, \"Cyclonic bad OW\", \"lightblue\"),\n):\n    x, y = d.bins_stat(**kw_a)\n    ax_amp.plot(x * 100, y, label=label, color=color)\n    ax_amp_c.plot(\n        x * 100, y.cumsum() / y.sum(), label=label, color=color, ls=\"-.\", lw=0.5\n    )\n    x, y = d.bins_stat(**kw_r)\n    ax_rad.plot(x * 1e-3, y, label=label, color=color)\n    ax_rad_c.plot(\n        x * 1e-3, y.cumsum() / y.sum(), label=label, color=color, ls=\"-.\", lw=0.5\n    )\n\nax_amp.set_xlim(0, 12.5), ax_amp.grid(), ax_amp.set_ylim(0), ax_amp.legend()\nax_rad.set_xlim(0, 120), ax_rad.grid(), ax_rad.set_ylim(0)\nax_amp.set_xlabel(\"Amplitude (cm)\"), ax_amp.set_ylabel(\"Nb eddies\")\nax_rad.set_xlabel(\"Speed radius (km)\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}